from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_community.document_loaders import TextLoader
from langchain_community.document_loaders import UnstructuredExcelLoader
from langchain_community.document_loaders import WebBaseLoader
from langchain_unstructured import UnstructuredLoader
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain_community.document_loaders import UnstructuredImageLoader
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import CSVLoader
from langchain_community.document_loaders import UnstructuredCSVLoader
from langchain_core.documents import Document
from unittest import result
from langchain_core.documents.base import Document
from logging import Logger
import os
import logging
import json
import base64
from openai import OpenAI
from constant import *
from utils.pdf_unstrucured_loader_util import PdfUnstructuredLoaderUtil
import pymysql
from sqlalchemy import create_engine, text
import pandas as pd
from urllib.parse import quote_plus as urlquote

from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.config.parser import ConfigParser
from marker.output import text_from_rendered, save_output

logger: Logger = logging.getLogger(__name__)

class LoadService:
    def __init__(self):
        pass

    '''
    文档加载类:
    Document:
        metadata:文档相关的元信息，例如文档的来源路径、作者、日期等 
        page_content:文档内容
    '''

    # 加载文本文件
    #load_type:
    # 1 langchain_text_load
    def load_text(self,file_path: str,load_type:str):
        loader = TextLoader(file_path)
        docs = loader.load()
        print(docs)
        return docs;

    
    # 加载md文件
    #load_type:
    #1 unstructured_markdown_load
    #2 file_markdown_load
    def load_markdown(self,file_path: str,load_type:str):
        if load_type == "unstructured_markdown_load":
            loader = UnstructuredMarkdownLoader(file_path, mode="elements")
            docs = loader.load()
            '''
            docs = self._get_docs_from_unstructured(docs)
            for item in docs:
                # 检查当前项是否为包含父子元素的元组
                if isinstance(item, tuple):
                    parent, child = item
                    print(f'父元素 - {parent.metadata["category"]}: {parent.page_content}')
                    print(f'子元素 - {child.metadata["category"]}: {child.page_content}')
                else:
                    print(f'{item.metadata["category"]}: {item.page_content}')
                    print("-" * 80)
            '''
            return docs;
        elif load_type == "file_markdown_load":
            with open(file_path, 'r', encoding='utf-8') as file:
                markdown_text = file.read()
            return [Document(
                page_content=markdown_text,
                metadata={"file_type": "markdown"}
            )]

    #加载图片文字
    #load_type:
    #1 unstructured_image_load
    #2 deepseek_image_load
    def load_image(self,file_path: str,load_type:str):
        if load_type == 'unstructured_image_load':
            loader = UnstructuredImageLoader(file_path)
            docs = loader.load()
            for doc in docs:
                print(f"{doc.metadata}\n{doc.page_content.strip()}")
            return docs
        elif load_type == 'deepseek_image_load':
            docs = self._deepseek_image_load(file_path)
            for doc in docs:
                print(f"{doc.metadata}\n{doc.page_content.strip()}")
            return docs

    
    #大模型分析图片
    def _deepseek_image_load(self,file_path: str):
        try:
            # 获取 DeepSeek API 密钥
            api_key = DEEPSEEK_API_KEY
            
            # 配置 DeepSeek 客户端
            client = OpenAI(
                api_key=api_key,
                base_url=DEEPSEEK_API_URL
            )
            with open(file_path, "rb") as image_file:
                base64_image = base64.b64encode(image_file.read()).decode('utf-8')
            message_json = json.dumps([
                {"type": "text", "text": "请详细描述这张图片的内容，并提取出图片中的文字信息"},
                {"type": "image", "image": {"data": base64_image, "format": "base64"}},
            ])
            response = client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {
                        "role": "user",
                        "content": message_json
                    }
                ],
                max_tokens=300
            )
            result: str | None = response.choices[0].message.content
            doc = Document(
                page_content=result if result else "",
                metadata={"source": file_path}
            )
            return [doc]
        except Exception as e:
            logger.error(f"DeepSeek 图像加载失败: {str(e)}")
            # 返回空文档而不是抛出异常
            return [Document(
                page_content=f"图像加载失败: {str(e)}",
                metadata={"source": file_path, "error": True}
            )]

    # 加载pdf文件
    #load_type:
    #1 pypdf_pdf_load
    #2 unstructured_pdf_load
    #3 marker_pdf_load
    def load_pdf(self,file_path: str,load_type:str):
        if load_type == 'pypdf_pdf_load':
            loader = PyPDFLoader(file_path)
            docs = loader.load()
            for doc in docs:
                print(f"{doc.metadata}\n{doc.page_content.strip()}")
            return docs
        elif load_type == 'unstructured_pdf_load':
            pdfUnstructuredLoaderUtil: PdfUnstructuredLoaderUtil = PdfUnstructuredLoaderUtil()
            return pdfUnstructuredLoaderUtil.load_pdf(file_path)
        elif load_type == 'marker_pdf_load':
            '''
            config = {
            "output_format": "markdown",  # 控制输出类型，可选值为："json"、"html"、"markdown"，默认是 "markdown"
            "ADDITIONAL_KEY": "VALUE",
            "format_lines": True,  # 等价于命令行 --format_lines，开启后会尝试对段落进行多行合并为一段，利于语义理解
            "force_ocr": True,  # 等价于 --force_ocr，强制对所有页面使用 OCR，即使页面中已有可提取的文字（如扫描件、图片）
            "use_llm": False,  # 等价于 --use_llm，开启后会使用 LLM 对提取的文本进行处理

            "llm_service": "marker.services.openai.OpenAIService",
            "openai_api_key": "your_api_key",
            "openai_base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",  # 假设为阿里云百炼平台
            "openai_model": "qwen-plus-latest",
            }
            '''

            config = {
                "output_format": "markdown",  # 控制输出类型，可选值为："json"、"html"、"markdown"，默认是 "markdown"
                "ADDITIONAL_KEY": "VALUE"
                }
            config_parser = ConfigParser(config)
            
            converter = PdfConverter(
                config=config_parser.generate_config_dict(),
                artifact_dict=create_model_dict(),
                processor_list=config_parser.get_processors(),
                renderer=config_parser.get_renderer()
            )
            rendered = converter(filepath=file_path)
            output_dir = file_path[:str.rindex(file_path, ".pdf")]
            os.makedirs(output_dir, exist_ok=True)
            # 文件标识名(一般为.pdf前面的名字，不包括父路径)
            fname_base = output_dir[str.rindex(output_dir, "/")+1:]
            save_output(rendered, output_dir, fname_base)
            # 保存为markdown后读取为文本并返回
            with open(output_dir+"/"+fname_base+".md", 'r', encoding='utf-8') as file:
                markdown_text = file.read()
            return [Document(
                page_content=markdown_text,
                metadata={"file_type": "markdown"}
            )]

            
    # 加载excel文件
    #load_type:
    def load_excel(self,file_path: str,load_type:str):
        loader = UnstructuredExcelLoader(file_path=file_path)
        data = loader.load()
        for record in data:
            print(record)
        return data


    # 加载csv文件
    #load_type:
    #1 csv_loader
    #2 unstructured_csv_loader
    def load_csv(self,file_path: str,load_type:str):
        if load_type == "csv_loader":
            loader = CSVLoader(file_path=file_path)
            data = loader.load()
            for record in data:
                print(record)
            return data
        elif load_type == "unstructured_csv_loader":
            loader = UnstructuredCSVLoader(file_path=file_path)
            data = loader.load()
            for record in data:
                print(record)
            return data


    # 加载mysql数据
    #load_type:    
    #1 pymysql_loader
    #2 sqlalchemy_loader
    def load_mysql(self,load_type: str,load_sql: str):
        if load_type == "pymysql_loader":
            try:
                connection = pymysql.connect(
                    host= mysql_ip,
                    user= mysql_username,
                    password= mysql_password,
                    database= mysql_database,
                    port= mysql_port
                )

                with connection.cursor() as cursor:
                    cursor.execute("SELECT * FROM sys_user")
                    for row in cursor.fetchall():
                        print(row)

                connection.close()

            except Exception as e:
                logger.error("数据库连接失败:", e)
        elif load_type == "sqlalchemy_loader":
            # 确保使用 pymysql 作为驱动
            engine = create_engine(f"mysql+pymysql://{mysql_username}:{urlquote(mysql_password)}@{mysql_ip}:{mysql_port}/{mysql_database}")
            # 测试连接
            try:
                with engine.connect() as connection:
                    # 使用 text() 函数包装 SQL 语句
                    result = connection.execute(text("SELECT * FROM sys_user"))
                    rows = result.fetchall()
                    df = pd.DataFrame(rows)
                    print("查询结果：")
                    print(df)
            except Exception as e:
                logger.error("数据库连接失败:", e)

    
    # 加载web网页
    #load_type: 
    # 1 langchain_html_load
    # 2 unstructured_html_load
    def load_html(self,web_url: str,load_type:str):
        if load_type == 'langchain_html_load' :
            loader = WebBaseLoader(web_url)
            docs = []
            docs = loader.load()
            for doc in docs:
                print(f"{doc.metadata}\n{doc.page_content.strip()}")
            return docs
        elif load_type == 'unstructured_html_load' :
            loader = UnstructuredLoader(web_url=web_url)
            docs = []
            docs: list[Document] = loader.load()
            '''
            docs = self._get_docs_from_unstructured(docs)
            for item in docs:
                # 检查当前项是否为包含父子元素的元组
                if isinstance(item, tuple):
                    parent, child = item
                    print(f'父元素 - {parent.metadata["category"]}: {parent.page_content}')
                    print(f'子元素 - {child.metadata["category"]}: {child.page_content}')
                else:
                    print(f'{item.metadata["category"]}: {item.page_content}')
                    print("-" * 80)
            '''
            return docs
    
    #将网页内容转化为包含父子关系的文档list
    def _get_docs_from_unstructured(self, docs: list[Document]):
        setup_docs= []
        parent_id = None # 初始化 parent_id
        current_parent = None # 用于存储当前父元素
        for doc in docs:
            # 检查是否是 Title 或 Table
            if doc.metadata["category"] == "Title" or doc.metadata["category"] =="Table":
                parent_id = doc.metadata["element_id"]
                current_parent = doc # 更新当前父元素
                setup_docs.append(doc)
            elif doc.metadata.get("parent_id") == parent_id:
                setup_docs.append((current_parent, doc)) # 将父元素和子元素一起存储
        return setup_docs
